The alias mxfp4_moe_forward_triton as mxfp4_moe_forward_fallback is highly misleading. It imports the optimized Triton implementation but names it as if it were the fallback. This is especially confusing given that a separate PyTorch fallback implementation exists in mxfp4_moe_fallback.py.

This can cause significant confusion during future maintenance. Please rename the alias to accurately reflect the imported function, for example, by simply importing mxfp4_moe_forward_triton and using that name directly.

Please address this comment

Done — now imports mxfp4_moe_forward_triton directly without aliasing.

The assert messages have been changed from descriptive explanations to simply "TODO". The previous implementation also had these asserts, but the original messages explained why these values were hardcoded (e.g., "torch reference impl hardcodes DSV4 indexer head_dim=128"). Please consider restoring the more descriptive messages to improve code maintainability.

dito

Done — restored descriptive assert messages.

In flash_mla_with_kvcache_entrypoint, the backend selection logic for SM120 is confusing. The function takes a backend parameter (from SGLANG_HACK_FLASHMLA_BACKEND), but this parameter is ignored when _is_sm120 is true. Instead, the choice between the Triton kernel and the PyTorch fallback is controlled by a separate environment variable, SGLANG_SM120_TRITON_FLASHMLA.

To improve clarity, it would be better to unify this control flow. Consider using the backend parameter to also control the implementation path on SM120, for instance, by checking for values like "triton" or "torch".

dito

Done — the SGLANG_HACK_FLASHMLA_BACKEND env var has been removed. On SM120 the backend is now hardcoded to "kernel" (Triton path), with SGLANG_SM120_TRITON_FLASHMLA env var as the only escape hatch to fall back to PyTorch attention if needed for debugging. This simplifies the control flow as the gemini-bot suggested.

Rename it to flash_mla_with_kvcache_sm120. Also, please move this function to a standalone file. It shouldn't be put in the fallback file

Created flash_mla_sm120.py as the standalone file with the function renamed to flash_mla_with_kvcache_sm120. Deleted the old flash_mla_sm120_fallback.py.

Does here need add _is_sm120() ?

I think is_cuda covers SM120

I mean and not _is_sm120

Done — added and not _is_sm120 guards at all DeepGEMM/tilelang paths in nsa_backend.py.

topk_ids can be -1 for padded/filtered tokens, but this path passes them directly as Triton expert ids. Could we use safe_expert_ids = topk_ids.clamp_min(0) for loads and set the invalid slots' output to zero?? The PyTorch fallback already skips eid_val < 0

Handled — see lines 361-363: flat_expert_ids_raw.clamp(min=0) for safe indexing, plus line 441+: invalid slots are zeroed out after the kernel. The Triton kernel itself runs on clamped-to-0 expert IDs (safe for loads), and the output is masked to zero post-kernel.

For SM120, should this CP ragged path explicitly raise NotImplementedError instead of falling through to deep_gemm.fp8_mqa_logits, since DeepGEMM is unsupported here?

Good catch. Will add if _is_sm120: raise NotImplementedError("CP ragged indexer not supported on SM120")

Currently, there is a little bug on old CUTLASS version if the kernel is CUTLASS. But can you elaborate on PDL not working?

PDL is now re-enabled on SM120.

Don't modify this function inplace. Please open a new function for fp8_paged_mqa for sm120 and put your implementation there

Done. Restored the original fp8_paged_mqa_logits_torch and created a separate fp8_paged_mqa_logits_torch_sm120 with the vectorized (no .item()) implementation. The dispatch in C4IndexerBackendMixin selects the SM120 variant when is_sm120_supported().

It has already been set in server_args.py

Removed. The SM120 marlin auto-select is now only in server_args.py.

dsa_prefill_backend and dsa_decode_backend are for glm-5 and dpsk v3.2. REmove their settings

Removed both SM120 branches from _set_default_dsa_backends().

Move these tests to main function of sm120_mqa_fallback.py and sm120_mqa_triton.py. Don't register them as CI tests

Removed the test file and CI registration. The source files (sm120_mqa_fallback.py, sm120_mqa_triton.py) under dsa/ were also removed per comment #6 (dead code for V4), so there is no target to move tests into.

instead of adding these is_sm120_supported() in the if condition, can we auto set these environs to false when we detect sm120

Done. SM120 env auto-set is now in server_args.py under the DeepseekV4ForCausalLM block:

envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
envs.SGLANG_OPT_USE_TOPK_V2.set(False)
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False)
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.set(True)

Removed all is_sm120_supported() guards from deepseek_v4.py. Also changed post_load_weights to use envs.SGLANG_OPT_FP8_WO_A_GEMM.get() instead of the module-level cached _FP8_WO_A_GEMM, since the env is set after module import.

This file is unrelated to dpsk v4. Don't modify it

Reverted all changes.

Use sm_version == 120 here instead

Changed to if sm_version == 120: return False.

Don't modify in this way. Please only use the newly created entrypoint when it's sm120

Restructured: SM120 branch calls flash_mla_with_kvcache_sm120 from the new standalone file; the original flash_mla.flash_mla_with_kvcache(...) call is preserved unchanged in the else branch.

Remove it if it's not used anywhere

Confirmed no imports. Deleted.

Rename it to flash_mla_with_kvcache_sm120. Also, please move this function to a standalone file. It shouldn't be put in the fallback file

pls revert this unnecessary change

Done — reverted in cd44380.